跳到主要内容

PyTorch TensorBoard 使用

TensorBoard 是什么？​

TensorBoard 是 TensorFlow 提供的一个可视化工具，主要用于可视化神经网络的训练过程、结构和其它各种指标。不过，TensorBoard 不仅限于 TensorFlow，它与 PyTorch 也有很好的集成。 为了将数字列表转换为 Markdown 子标题，我们可以使用 ###。这是怎么做的：

使用 TensorBoard 与 PyTorch，你可以：

  • 可视化模型结构。
  • 跟踪和可视化模型的训练指标，如损失和准确率。
  • 查看梯度的分布和趋势。
  • 查看权重、偏置和其他张量的分布和变化。
  • 生成图片、音频和文本的可视化。

要在 PyTorch 中使用 TensorBoard，以下是基本步骤：

安装​

pip install tensorboard

导入必要的库

from torch.utils.tensorboard import SummaryWriter

创建 SummaryWriter​

这个对象是你与 TensorBoard 交互的主要方式。

writer = SummaryWriter('runs/experiment_name')

记录信息​

使用 writer 对象来记录想要可视化的信息。

for epoch in range(num_epochs):
    # ... training loop ...
    writer.add_scalar('Loss/train', loss, epoch)
    writer.add_scalar('Accuracy/train', accuracy, epoch)

启动 TensorBoard​

在命令行或终端中运行：

tensorboard --logdir=runs

之后，你可以在浏览器中打开 http://localhost:6006/ 来查看 TensorBoard 的界面。

每次启动记录新的实验​

tensorboard 会在一个目录下持续记录，如果想要记录新的实验，如果不重新创建目录去记录，就会出现下面这样奇怪的图

解决办法就是使用时间戳随机生成一个目录名称

from datetime import datetime

now = datetime.now()
logdir = "runs/" + now.strftime("%Y%m%d-%H%M%S") + ""
writer = SummaryWriter(logdir)

TensorBoard 记录模型的结构​

要使用 TensorBoard 记录模型的结构，可以使用 TensorBoard 的 add_graph 方法。这个方法可以捕获模型的计算图并在 TensorBoard 中显示它。以下是如何做到这一点的步骤：

创建 SummaryWriter 对象

writer = SummaryWriter('runs/experiment_name')

记录模型的计算图

为了记录模型的结构，你需要模型的一个实例和一个代表输入的张量。

import torch
import torchvision.models as models

# 获取一个预训练的 resnet18 模型
model = models.resnet18(pretrained=True)
# 创建一个随机输入张量，假设批量大小为 1，通道数为 3，高度和宽度都为 224
input_tensor = torch.rand(1, 3, 224, 224)

# 使用 add_graph 方法记录模型的结构
writer.add_graph(model, input_tensor)

add_graph 函数的主要用途是在 TensorBoard 中可视化模型的计算图。为了实现这一点，该函数需要两个主要输入：模型本身和一个输入张量（即模型的输入数据）。

输入张量的作用有两点：

  1. 形状推断：通过将输入张量传递给模型，TensorBoard 可以捕获模型的整个计算图，包括每一层的输入和输出的形状。这有助于为 TensorBoard 用户呈现每一层的维度和大小。

  2. 正向传播：输入张量确实会通过模型进行一次前向传播，但实际上不会进行训练（即不会进行反向传播）。这次前向传播的目的是捕获模型的完整计算图。

例如，对于一个典型的卷积神经网络，可能想知道在每一层卷积、池化或全连接层之后，特征图的形状是什么。提供输入张量可以帮助捕获这些信息，并在 TensorBoard 中可视化模型的层次结构和每一层的输出形状。

在实际的训练中，这里可以选取数据集的一部分作为输入张量，然后将其传递给模型。如下代码所示

# 训练模型
num_epochs = 5  # 这只是一个示例值，您可以根据需要进行调整
for epoch in range(num_epochs):
    total_loss = 0.0
    for i, (inputs, labels) in enumerate(trainloader):
        inputs, labels = inputs.to(device), labels.to(device)
        
        # 清零梯度
        optimizer.zero_grad()
        
        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs.view(-1, len(nums)), labels.view(-1))
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

        # 每100批次，记录损失到TensorBoard
        if (i + 1) % 100 == 0:
            writer.add_scalar('Training Loss', total_loss / (i + 1), epoch * len(trainloader) + i)
    
    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {total_loss / (i + 1):.4f}")

# 在 TensorBoard 中记录模型的结构
inputs, _ = next(iter(trainloader))
inputs = inputs.to(device)
writer.add_graph(model, inputs)

# 关闭 SummaryWriter
writer.close()

References​